package com.example.ymm.config;

import com.alibaba.fastjson.JSONObject;
import com.example.ymm.annotation.RequestLimiter;
import com.example.ymm.model.CommonResult;
import com.example.ymm.util.StringUtils;
import com.example.ymm.util.UtilValidate;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Configuration
@Slf4j
public class RateLimiterInterceptor implements HandlerInterceptor {

	private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();
	/**
	 * 这个方法的主要作用是用于清理资源的，当然这个方法也只能在当前这个Interceptor的preHandle方法的返回值为true时才会执行。
	 */
	@Override
	public void afterCompletion(HttpServletRequest arg0, HttpServletResponse arg1, Object arg2, Exception arg3)
	{
		// TODO Auto-generated method stub

	}

	/**
	 * 在Controller的方法调用之后执行
	 */
	@Override
	public void postHandle(HttpServletRequest arg0, HttpServletResponse response, Object arg2, ModelAndView arg3)
			throws IOException{

	}


	@Override
	public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
		CommonResult<JSONObject> jsonObject = CommonResult.failed("操作太频繁，请稍后再试!");
		try {
			if (handler instanceof HandlerMethod) {
				HandlerMethod handlerMethod = (HandlerMethod) handler;
				RequestLimiter rateLimit = handlerMethod.getMethodAnnotation(RequestLimiter.class);
				if(UtilValidate.isNotEmpty(rateLimit)){
					if(UtilValidate.isNotEmpty(rateLimit.msg())){
						jsonObject.setMessage(rateLimit.msg());
					}
					//判断是否有注解
					if (rateLimit != null) {
						// 获取请求url
						String ip = StringUtils.getIp(request);
						String url = request.getRequestURI();
						String key=ip+":"+url;
						RateLimiter rateLimiter;
						// 判断map集合中是否有创建好的令牌桶
						if (!rateLimiterMap.containsKey(key)) {
							// 创建令牌桶,以n r/s往桶中放入令牌
							rateLimiter = RateLimiter.create(rateLimit.QPS());
							rateLimiterMap.put(key, rateLimiter);
						}
						rateLimiter = rateLimiterMap.get(key);
						// 获取令牌
						boolean acquire = rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.timeunit());
						if (acquire) {
							//获取令牌成功
							return true;
						} else {
							log.warn("请求被限流,url:{}", request.getServletPath());
							makeResult(response, renderJson(jsonObject));
							return false;
						}
					}
				}

			}
			return true;
		} catch (Exception var6) {
			var6.printStackTrace();
			makeResult(response, renderJson(jsonObject));
			return false;
		}
	}

	private void makeResult(HttpServletResponse response, JSONObject jo) {
		response.setContentType("application/json; charset=utf-8");
		response.setCharacterEncoding("UTF-8");
		try (PrintWriter out = response.getWriter()) {
			out.append(jo.toJSONString());
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	private JSONObject renderJson(Object o) {
		return JSONObject.parseObject(JSONObject.toJSONString(o));
	}

}
